# -*- coding: utf-8 -*-
'''
@File    :   dialogueAgnet.py
@Author  :   一力辉 
'''

from ..builder import AGENTS
from ..agentBase import BaseAgent
from typing import List, Optional, Union
from digitalHuman.utils import logger
from digitalHuman.utils import AudioMessage, TextMessage
from digitalHuman.engine.engineBase import BaseEngine
from dashscope import Application
from http import HTTPStatus
import re
import httpx
import json

__all__ = ["AliAgent"]


@AGENTS.register("AliAgent")
class AliAgent(BaseAgent):

    def checkKeys(self) -> List[str]:
        return ["KEY", "APPID"]
    
    async def run(
        self, 
        input: Union[TextMessage, AudioMessage], 
        streaming: bool,
        **kwargs
    ):
        try: 
            if isinstance(input, AudioMessage):
                raise RuntimeError("AliAgent does not support AudioMessage input")
            
            # responses = Application.call(app_id=self.cfg.APPID,#YOUR-APPID
            #                      prompt=input.data,
            #                      api_key=self.cfg.KEY, #YOUR-API-KEY
            #                      stream=True,
            #                      incremental_output=True
            #                      ) 

            # for response in responses:
            #     if response.status_code != HTTPStatus.OK:
            #         print('request_id=%s, code=%s, message=%s\n' % (
            #             response.request_id, response.status_code, response.message))            
            #     else:
            #         print('output=%s, usage=%s\n' % (response.output.text, response.usage))
            #         yield bytes(response.output.text, encoding='utf-8')

            # async with Application.call(app_id=self.cfg.APPID,  
            #                      prompt=input.data,
            #                      api_key=self.cfg.KEY, #YOUR-API-KEY
            #                      stream=True,
            #                      incremental_output=True
            #                      ) as responeses:

            #     async for response in responeses:
            #         if response.status_code != HTTPStatus.OK:
            #             print('request_id=%s, code=%s, message=%s\n' % (
            #                 response.request_id, response.status_code, response.message))            
            #         else:
            #             print('output=%s, usage=%s\n' % (response.output.text, response.usage))
            #             yield bytes(response.output.text, encoding='utf-8')

            headers = {
                    'Content-Type': 'application/json',
                    'Authorization': f'Bearer {self.cfg.KEY}',
                    'X-DashScope-SSE': 'enable'
                }
            payload = {
                "input": {"prompt": input.data},
                "parameters":  {"incremental_output": "true"},
                "debug": {}
            }

            pattern = re.compile(r'data:\s*({.*})')
            client = httpx.AsyncClient(headers=headers)

            async with client.stream('POST', f'https://dashscope.aliyuncs.com/api/v1/apps/{self.cfg.APPID}/completion', headers=headers, json=payload) as response:
                async for chunk in response.aiter_bytes():
                    chunkStr = chunk.decode('utf-8').strip()
                    chunkData = pattern.search(chunkStr)
                    if not chunkData: 
                        print(f"[AGENT] Engine return truncated data: {chunkData}")
                        continue
                    chunkData = chunkData.group(1)

                    try:
                        data = json.loads(chunkData)
                        if data['output']['finish_reason'] == "stop":
                            logger.info(f"AliAgent输出结束。最后一句{data['output']['text']}")
                            yield {"data" : {'text' :data['output']['text'],"finish_status" : "true"}}
                        else:
                            logger.info(f"AliAgent输出： {data['output']['text']}")
                            yield {"data": {'text': data['output']['text'], "finish_status": "false"}}
                    except Exception as e:
                        logger.info(f"[AGENT] Engine run failed: {e}")
                        yield {"data" : {'text' : "内部错误，请检查dify信息。","finish_status" : "true"}}
        except Exception as e:
            logger.error(f"[AGENT] Engine run failed: {e}")
            yield bytes("", encoding='utf-8')